# -*- coding: utf-8 -*-
"""
@File: MyGaitNet.py
@Time: 2021/12/19 13:26
@Author: 鹄望潇湘
@desc: 

"""
import torch.nn as nn
from backbone.CoAtNet import CoAtNet
import torch


class CoAtGait(nn.Module):
    def __init__(self, input_channel, image_size):
        super(CoAtGait, self).__init__()
        self.coatNet = CoAtNet(input_channel, image_size)
        self.fc1 = nn.Sequential(
            nn.Linear(768 * 4 * 4, 2048)
        )
        self.__initialize_weights()
        pass

    def __initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight.data)
                nn.init.constant_(m.bias.data, 0.0)

    def forward(self, x):
        x = self.coatNet(x)
        x = torch.flatten(x, start_dim=1)  #
        x = self.fc1(x)
        return x
